'''
Created on Apr 28, 2011

@author: jlucasgl
'''

# from yacf.MiddleEnd.Loop.Analysis import ParametrizeLoopTool
from yacf.Backends.Common.Mutators.AbstractMutator import AbstractMutator

from yacf.Tools.Tree import ReplaceTool, getContainerAttribute, InsertTool
from yacf.Frontend.Shortcuts import getCurrentLanguageAst
from yacf.Backends.Common.Visitors.GenericVisitors import GenericFilterVisitor,\
    ArrayRefVisitor, AnyIDVisitor
from yacf.Tools.Tree import NodeNotFound
from yacf.MiddleEnd.DataAnalysis.LlcScope import ArrayIndexAnalysis
from yacf.Frontend.C99 import c99_ast
from yacf.MiddleEnd.Loop.Analysis import ParametrizeLoopTool
from yacf.Tools.Debug import DGraphDebug
from yacf.Frontend.InternalRepr import AstToIR

# self._aia.extractIndexVars(aref)
class InvariantRefFilter(GenericFilterVisitor):
    def _extractIndexVars(self, node):
        if not isinstance(node, c99_ast.ArrayRef):
            raise NotImplemented
        # For each var referenced inside subscript
        idVisitor = AnyIDVisitor()
        idList = [ elem for elem in idVisitor.iterate(node.subscript)]
        return idList
    
    def _checkInvariant(self, aref, forStmt):
        try:
            params = ParametrizeLoopTool().apply(forStmt)
            lvs = self._st.lookUp(params['loop_variable'])
            self._indexSymbols = [self._st.lookUp(iVar) for iVar in self._extractIndexVars(aref)]
            if lvs in self._indexSymbols:
                return False
            
            for var in self._indexSymbols:
                # Case there is a dependency bw a var inside index and the loop variable 
                #  transformation is not feasible
                if self._dGraph.checkDependency(self._dGraph[lvs], self._dGraph[var]):
                    # DGraphDebug().apply(self._dGraph)
                    print " Aref is loop dependent, cannot optimize"
                    return False
            else:
                # print " Statement is loop-independent, transform "
                return True
        except NodeNotFound:
            raise NotImplemented
    
    
    def __init__(self, st, dGraph):
        self._for = 0
        self._for_stmt = None
        self._st = st
        self._dGraph = dGraph
        # Create graph
        # self._aia = ArrayIndexAnalysis(dGraph, st)
        # Get index Variables
        # self._aia.visit(subtree)
        self._analyzed = set()
        self._relatedForStmt = dict()
        
        def condition_func(aref):
            # if isinstance(aref, c99_ast.ArrayRef) and (self._for == 1) and self._for_stmt:
            if isinstance(aref, c99_ast.ArrayRef)  and self._for_stmt:
                # check invariant of aref related to self._for_stmt
                if self._checkInvariant(aref, self._for_stmt) \
                        and not aref in self._analyzed \
                        and not isinstance(aref.parent, c99_ast.ArrayRef):
                    print " Aref: " + str(aref) + " is marked as potential invariant array access"
                    self._analyzed.add(aref)
                    self._relatedForStmt[aref.name] = self._for_stmt
                    return True
            return False
        super(InvariantRefFilter, self).__init__(condition_func = condition_func,
                                                    prev_brother = None,)

    def getForStmt(self, aref_name):
        return self._relatedForStmt[aref_name]
    
    def visit_For(self, node, prev, offset = 1, ignore = []):
        self._for += 1
        prev = self._for_stmt
        self._for_stmt = node
        rvalue =  self.generic_visit(node, offset, ignore)
        self._for_stmt = prev
        self._for -= 1
        return rvalue



class LoopInvariantMutator(AbstractMutator):
    def __init__(self, st, dGraph):
        self._st = st
        self._dGraph = dGraph
        self._ivf = InvariantRefFilter(self._st, self._dGraph)
        self._mutatorCount = 0
 
    def filter(self, subtree):
        """  """
        return self._ivf.apply(subtree)

    def filter_iterator(self, subtree):
        """ Search for invariant Ref accesses inside a particular subtree """
        return self._ivf.iterate(subtree)

    def fast_filter(self, subtree):
        return self._ivf.iterate(subtree)
    
    def _check_equivalent_subscript(self, subtree_a, subtree_b):
        # TODO: CHANGE THIS!!!!
        return (str(subtree_a) == str(subtree_b))

    
    def mutatorFunction(self, aref):
        self._mutatorCount += 1
        innerARef = aref
        invariantSymbol = self._st.lookUp(innerARef)
        invariantLoop = self._ivf.getForStmt(innerARef.name)
        
        print " Array access " + str(aref) + " is invariant to loop iterating with (" + str(invariantLoop.next) +")"
        
        
        # Create new private/register var before  invariantStmt
        # name, quals, storage, funcspec, type, init, bitsize
        # TODO: Add numbering to this variable
        privateVarName = "__invariant_tmp_" + str(self._mutatorCount)
        idNode = c99_ast.ID(name = privateVarName)
        
        # type __invariant_tmp; 
        newVar_node = c99_ast.Decl(name = privateVarName, quals = [], 
                                    storage = [], 
                                    funcspec = [], 
                                    bitsize = None, 
                                    type = invariantSymbol.type.type, 
                                    init = None, 
                                    parent = invariantLoop.parent)
        #  __invariant_tmp = array[position];
        update1_node = c99_ast.Assignment(op = '=', 
                                          lvalue = idNode,
                                          rvalue = AstToIR.deepcopySubtree(innerARef),  
                                          parent = invariantLoop.parent
                                          )
        # array[position] =  __invariant_tmp ;
        update2_node = c99_ast.Assignment(op = '=', 
                                          lvalue = AstToIR.deepcopySubtree(innerARef),
                                          rvalue = idNode,  
                                          parent = invariantLoop.parent
                                          )
        # Insert declaration
        InsertTool(subtree = newVar_node, position = "before", next = invariantLoop).apply(invariantLoop.parent, getContainerAttribute(invariantLoop))
        # Insert tmp update
        InsertTool(subtree = update1_node, position = "before", next = invariantLoop).apply(invariantLoop.parent, getContainerAttribute(invariantLoop))
        # Insert value update from tmp
        InsertTool(subtree = update2_node, position = "after", prev = invariantLoop).apply(invariantLoop.parent, getContainerAttribute(invariantLoop))
        # Replace all ocurrences of invariantStmt by newVar inside invariantLoop
        for aref in ArrayRefVisitor().iterate(invariantLoop.stmt):
            if self._st.lookUp(aref.name) == invariantSymbol and self._check_equivalent_subscript(innerARef.subscript, aref.subscript):
                ReplaceTool(new_node = idNode, old_node = aref).apply(aref.parent, getContainerAttribute(aref))

        print " -> Extracted an replaced by private var " + str(privateVarName)

